It looks like the package is not yet on PyPI. Do we want to hold off on the PR until the package is published?

Yes, the package is not on PyPI yet. I have updated the error message in wpi_engine.py to remove the PyPI reference and only provide instructions on how to install it directly from the WPI source (cd weight-propagation-interface/consumer/wpi_client && pip install -e .).

Just FYI we are trying to clean up the usage of mode in the weight transfer APIs. #37476 is renaming this to send_mode to be explicit. (This is a pretty clunky part of the code that we want to remove long term, but thought I'll mention the planned change short term)

Thanks for the heads-up! Makes sense to make it more explicit. I've kept it as mode for now in the example to match the current WPITrainerSendWeightsArgs definition in this branch.

I'll keep an eye on #37476 and will update it to send_mode once that lands.

@yangspirit send_mode has landed. Can we update the PR to be consistent?

Done. PTAL.

Just for my understanding, when would callers use this pre-initialized context vs not?

Callers should use the pre-initialized context when performing repeated weight transfers (like in an RLHF training loop). Initializing it once at the start avoids the overhead of staging the buffer and importing CUDA memory on every step.

They would omit it (relying on lazy init) for one-off transfers or simple testing where they want to avoid managing the context object and don't care about the initialization overhead.

Good to know, thanks!

should follow start, for loop, finish contract here.

Done.

Codex found a problem I feel like make a lot of sense:
The buffer is byte-addressable, but the receiver reinterprets each slice with .view(dtype=dtype) without copying it into a new storage. Therefore, the recorded offset must be divisible by the target dtype’s itemsize.
For example, a 6-byte FP16 tensor followed by an FP32 tensor places the FP32 tensor at offset 6. The receiver then attempts buffer[6:...].view(torch.float32), which fails because the underlying storage offset is not divisible by 4.
Please align offset before recording and copying each tensor, and include the resulting padding when calculating buffer_size_bytes, for example:
for name, tensor in iterator:
weight = tensor.detach().contiguous()
itemsize = weight.element_size()
offset = (offset + itemsize - 1) // itemsize * itemsize
nbytes = weight.nbytes
The example’s total_model_bytes = sum(...) must also use the same aligned-layout calculation; otherwise the aligned final tensor can exceed the allocated buffer.

Done

Start the WPI update session before sending

In this HTTP example, trainer_send_weights() posts to /update_weights, but the script never calls /start_weight_update first or /finish_weight_update afterward. The vLLM worker rejects /update_weights unless _weight_update_active was set by start_weight_update, so following this example fails before any WPI-loaded weights are applied; wrap this send with the same start/finish calls used by the NCCL and IPC HTTP examples.

Useful? React with 👍 / 👎.

Done

Remove unsupported argument from start_weight_update

Users copying this WPI snippet will get a TypeError before the transfer starts because the vLLM LLM.start_weight_update()/HTTP start API takes no is_checkpoint_format parameter; the backend already chooses layerwise reload internally. Document llm.start_weight_update() so the example matches the actual public API.

Useful? React with 👍 / 👎.

Done

Disable sharded WPI until updates are re-sliced

When callers set total_shards > 0 as documented for tensor-parallel WPI, this asks the driver to stage/open only a shard, but trainer_send_weights still packs one full checkpoint and receive_weights later interprets each mapped buffer using full-buffer offsets and full tensor shapes. In any real scatter run, workers will read the wrong byte ranges or run past their local shard instead of receiving the full tensors that load_weights expects, so the advertised sharded WPI path corrupts or fails weight updates; keep this in broadcast mode or send per-shard/rebased metadata.

Useful? React with 👍 / 👎.

Done.

I found a codex comment make sense to me, could you check?

stage_weight() above receives shard_index and total_shards, and WPI registers a sharded buffer under the effective ID <buffer_id>_shard. However, the returned context only preserves the base buffer_id, discarding the shard identity that is needed by later operations.
Consequently, trainer_send_weights() calls propagate(buffer_id=ctx.buffer_id) with weights, while the staged source resource is actually weights__shard_N, so the driver cannot find the buffer.
Please store shard_index/total_shards or the effective buffer ID in WPITrainerContext, and use that same resource identity for propagation and cleanup.

Done.

Is it ok that for every tp rank use the same buffer id?

Done

And claim_id is the same, seems it's related to shutdown and resource release

Done